# -*- coding: utf-8 -*-
"""
Created on Mon Jun  3 14:38:12 2019

@author: guanqun
Written for plotting RMCProfile outputs
"""
import matplotlib.pyplot as plt
from matplotlib import rcParams
rcParams['mathtext.default'] = 'regular'

def dr(filename,gap=1.5):
    
    r=[]
    dr_rmc=[]
    dr_exp=[]
    error_dr = []
    dr=open(filename)
    title =dr.readline()                          #pass the first line
    for lines in dr.readlines():
        r.append(float(lines.split(',')[0]))
        dr_rmc.append(float(lines.split(',')[1]))
        dr_exp.append(float(lines.split(',')[2]))
        error_dr.append(float(lines.split(',')[1])-float(lines.split(',')[2])-gap)
    dr.close()
    plt.figure(figsize=(6,4))
   # plt.title('RMC PDF refinement',weight = 'bold')
    plt.plot(r, dr_exp,'k.',label = 'EXP',markersize=2)
    plt.plot(r, dr_rmc, 'r-', label = 'RMC',linewidth=1.0)
    plt.plot(r, error_dr, label = 'residual',linewidth=0.5)
    plt.xlabel('r ( $\AA$ )')
    plt.ylabel('D(r)')  
    plt.legend(fontsize = 14)
    plt.tight_layout()
 #   plt.savefig('%s.pdf'%filename.replace('.csv',''))
    
def sq(filename):
    
    Q=[]
    sq_rmc=[]
    sq_exp=[]
    error_sq = []
    sq=open(filename)
    title =sq.readline()                          #pass the first line
    for lines in sq.readlines():
        Q.append(float(lines.split(',')[0]))
        sq_rmc.append(float(lines.split(',')[1]))
        sq_exp.append(float(lines.split(',')[2]))
        error_sq.append(float(lines.split(',')[1])-float(lines.split(',')[2])-1.5)
    sq.close()
    plt.figure(figsize=(6,4))
    #plt.title('RMC scattering data refinement',weight = 'bold')
    plt.plot(Q, sq_exp,'k.',label = 'EXP',markersize=2)
    plt.plot(Q, sq_rmc, 'r-', label = 'RMC',linewidth=1.0)
    plt.plot(Q, error_sq, label = 'error',linewidth=0.5)
    plt.xlabel('Q  ( $\AA^{-1}$ )')
    plt.ylabel('S(Q)')  
    plt.legend()
    plt.xlim(0,30)
    plt.tight_layout()
 #   plt.savefig('%s.pdf'%filename.replace('.csv',''))
def iq(filename):
    
    Q=[]
    iq_rmc=[]
    iq_exp=[]
    error_iq = []
    iq=open(filename)
    title =iq.readline()                          #pass the first line
    for lines in iq.readlines():
        Q.append(float(lines.split(',')[0]))
        iq_rmc.append(float(lines.split(',')[1]))
        iq_exp.append(float(lines.split(',')[2]))
        error_iq.append(float(lines.split(',')[1])-float(lines.split(',')[2])-1.2)
    iq.close()
    plt.figure(figsize=(6,4))
    #plt.title('RMC scattering data refinement',weight = 'bold')
    plt.plot(Q, iq_exp,'k.',label = 'EXP',markersize=2)
    plt.plot(Q, iq_rmc, 'r-', label = 'RMC',linewidth=1.0)
    plt.plot(Q, error_iq, label = 'residual',linewidth=0.5)
    plt.xlabel('Q  ( $\AA^{-1}$ )')
    plt.ylabel('i(Q)')  
    plt.xlim(0,30)
    plt.legend(fontsize = 14)
    plt.tight_layout()
  #  plt.savefig('%s.pdf'%filename.replace('.csv',''))
def bragg(filename,gap=2): 
    tof=[]
    bragg_rmc=[]
    bragg_exp=[]
    error_bragg = []
    bragg=open(filename)
    title =bragg.readline()                          #pass the first line
    for lines in bragg.readlines():
        tof.append(float(lines.split(',')[0]))
        bragg_rmc.append(float(lines.split(',')[2]))
        bragg_exp.append(float(lines.split(',')[1]))
        error_bragg.append(float(lines.split(',')[2])-float(lines.split(',')[1])-gap)
    bragg.close()
    plt.figure(figsize=(6,4))
    #plt.title('RMC bragg refinement',weight = 'bold')
    plt.plot(tof, bragg_exp,'k.',label = 'EXP',markersize=2)
    plt.plot(tof, bragg_rmc, 'r-', label = 'RMC',linewidth=1.0)
    plt.plot(tof, error_bragg, label = 'residual',linewidth=0.5)
    plt.xlabel( 'tof ($\mu$s)')
    plt.ylabel('Intensity')  
    plt.legend(fontsize = 14)
    plt.tight_layout()
  #  plt.savefig('%s.pdf'%filename.replace('.csv',''))
def partials(filename):
    partials=open(filename)
    r=[]
    bonds={}
    bond_name={}
    title =partials.readline() #pass the first line
    nlabel = len(title.split(','))
    for i in range (nlabel-1):
        bonds[i+1] = []
        bond_name[i+1] = title.split(',')[i+1]
    for lines in partials.readlines():
        r.append(float(lines.split(',')[0]))
        for i in range (nlabel-1):
            bonds[i+1].append(float(lines.split(',')[i+1]))        
    partials.close()   
    plt.figure()
    plt.title('RMC refined partial PDFs',weight = 'bold')
    for i in range (nlabel-1):
        plt.plot(r, bonds[i+1],label = bond_name[i+1],linewidth=1.0)
    plt.xlabel( 'r  ($\AA$)')
    plt.xlim(0,12)
    plt.ylabel(r'$g_{ij} (r)$')    
    plt.legend(fontsize = 14)   
    plt.tight_layout()
    
def read_dr(filename):
    r=[]
    dr_rmc=[]
    dr_exp=[]
    error_dr=[]
    dr=open(filename)
    title =dr.readline()                          #pass the first line
    for lines in dr.readlines():
        r.append(float(lines.split(',')[0]))
        dr_rmc.append(float(lines.split(',')[1]))
        dr_exp.append(float(lines.split(',')[2]))
        error_dr.append(float(lines.split(',')[1])-float(lines.split(',')[2])-2)
    dr.close()    
    return (r, dr_rmc, dr_exp, error_dr)

def read_bragg(filename):
    tof=[]
    bragg_rmc=[]
    bragg_exp=[]
    error_bragg = []
    bragg=open(filename)
    title =bragg.readline()                          #pass the first line
    for lines in bragg.readlines():
        tof.append(float(lines.split(',')[0]))
        bragg_rmc.append(float(lines.split(',')[2]))
        bragg_exp.append(float(lines.split(',')[1]))
        error_bragg.append(float(lines.split(',')[2])-float(lines.split(',')[1])-0.3)
    bragg.close()    
    return (tof, bragg_rmc, bragg_exp, error_bragg)

def read_sq(filename):
    Q=[]
    sq_rmc=[]
    sq_exp=[]
    error_sq = []
    sq=open(filename)
    title =sq.readline()                          #pass the first line
    for lines in sq.readlines():
        Q.append(float(lines.split(',')[0]))
        sq_rmc.append(float(lines.split(',')[1]))
        sq_exp.append(float(lines.split(',')[2]))
        error_sq.append(float(lines.split(',')[1])-float(lines.split(',')[2])-1.5)
    sq.close()
    return (Q, sq_rmc, sq_exp, error_sq)

def read_partials(filename):
    partials=open(filename)
    rp=[]
    bonds={}
    bond_name={}
    title =partials.readline() #pass the first line
    nlabel = len(title.split(','))
    for i in range (nlabel-1):
        bonds[i+1] = []
        bond_name[i+1] = title.split(',')[i+1]
    for lines in partials.readlines():
        rp.append(float(lines.split(',')[0]))
        for i in range (nlabel-1):
            bonds[i+1].append(float(lines.split(',')[i+1]))        
    partials.close()   
    return (rp,bonds,bond_name, nlabel)

'''
The crystal function plot all refine results excluding partial pdfs in a single plot for crystalline samples
It needs 3 file names as arguments
Currently the 3 files must follow the order of pdf, sq, bragg
'''
def crystal(file1, file2, file3): 
    Q=[]
    sq_rmc=[]
    sq_exp=[]
    error_sq=[]
    r=[]
    dr_rmc=[]
    dr_exp=[]
    error_dr=[]
    tof=[]
    bragg_rmc=[]
    bragg_exp=[]
    error_bragg = []
    Q, sq_rmc, sq_exp, error_sq = read_sq(file2)
    tof, bragg_rmc, bragg_exp, error_bragg = read_bragg(file3) 
    r, dr_rmc, dr_exp, error_dr = read_dr(file1)
    plt.figure()
    plt.subplot(2,2,3)
    plt.text(5000,1,'(b)')
    #plt.title('RMC bragg refinement',weight = 'bold')
    plt.plot(tof, bragg_exp,'k.',label = 'EXP')
    plt.plot(tof, bragg_rmc, 'r-', label = 'RMC',linewidth=1.0)
    plt.plot(tof, error_bragg, label = 'error',linewidth=0.5)
    plt.xlabel( 'tof ($\mu$s)')
    plt.ylabel('Intensity')  
    plt.legend()

    plt.subplot(2,2,4)
   # plt.title('RMC scattering data refinement',weight = 'bold')
    plt.text(0,2.5,'(c)')
    plt.plot(Q, sq_exp,'k.',label = 'EXP')
    plt.plot(Q, sq_rmc, 'r-', label = 'RMC',linewidth=1.0)
    plt.plot(Q, error_sq, label = 'error',linewidth=0.5)
    plt.xlabel('Q  ( $\AA^{-1}$ )')
    plt.ylabel('S(Q)')  
    plt.legend()

    plt.subplot(2,1,1)
    plt.text(-0.5,2.2,'(a)')
   # plt.title('RMC PDF refinement',weight = 'bold')
    plt.plot(r, dr_exp,'k.',label = 'EXP')
    plt.plot(r, dr_rmc, 'r-', label = 'RMC',linewidth=1.0)
    plt.plot(r, error_dr, label = 'error',linewidth=0.5)
    plt.xlabel('r ( $\AA$ )')
    plt.ylabel('D(r)')  
    plt.legend()
    plt.tight_layout()

'''
The alloutputs function plot all refine results including partial pdfs in a single plot
It needs 3 file names as arguments
Currently the 3 files must follow the order of pdf, sq, bragg, PDFpartials
'''
def alloutputs(file1, file2, file3, file4): 
    Q=[]
    sq_rmc=[]
    sq_exp=[]
    error_sq=[]
    r=[]
    dr_rmc=[]
    dr_exp=[]
    error_dr=[]
    tof=[]
    bragg_rmc=[]
    bragg_exp=[]
    error_bragg = []
    rp=[]
    bonds={}
    bond_name={}
    Q, sq_rmc, sq_exp, error_sq = read_sq(file2)
    tof, bragg_rmc, bragg_exp, error_bragg = read_bragg(file3) 
    r, dr_rmc, dr_exp, error_dr = read_dr(file1)
    rp, bonds, bond_name, nlabel = read_partials(file4)
    plt.figure()
    plt.subplot(2,2,3)
    plt.title('RMC bragg refinement',weight = 'bold')
    plt.plot(tof, bragg_exp,'k.',  markersize=2, label = 'EXP')
    plt.plot(tof, bragg_rmc, 'r-', label = 'RMC',linewidth=1.0)
    plt.plot(tof, error_bragg, label = 'error',linewidth=0.5)
    plt.xlabel( 'tof ($\mu$s)')
    plt.ylabel('Intensity')  
    plt.legend()

    plt.subplot(2,2,2)
    plt.title('RMC scattering data refinement',weight = 'bold')
    plt.plot(Q, sq_exp,'k.', markersize=2, label = 'EXP')
    plt.plot(Q, sq_rmc, 'r-', label = 'RMC',linewidth=1.0)
    plt.plot(Q, error_sq, label = 'error',linewidth=0.5)
    plt.xlabel('Q  ( $\AA^{-1}$ )')
    plt.ylabel('S(Q)')  
    plt.legend()

    plt.subplot(2,2,1)
    plt.title('RMC PDF refinement',weight = 'bold')
    plt.plot(r, dr_exp,'k.',markersize=2,label = 'EXP')
    plt.plot(r, dr_rmc, 'r-', label = 'RMC',linewidth=1.0)
    plt.plot(r, error_dr, label = 'error',linewidth=0.5)
    plt.xlabel('r ( $\AA$ )')
    plt.ylabel('D(r)')  
    plt.legend()
    
    plt.subplot(2,2,4)
    plt.title('RMC refined partial PDFs',weight = 'bold')
    for i in range (nlabel-1):
        plt.plot(rp, bonds[i+1],label = bond_name[i+1],linewidth=1.0)
    plt.xlabel( 'r  ($\AA$)')
    plt.xlim(0,12)
    plt.ylabel(r'$g_{ij} (r)$')    
    plt.legend()  
    
    plt.tight_layout()

'''
The glass function plot all refine results excluding partial pdfs in a single plot for glass samples
It needs 2 file names as arguments
Currently the 2 files must follow the order of pdf, sq
'''
def glass(file1, file2): 
    Q=[]
    sq_rmc=[]
    sq_exp=[]
    error_sq=[]
    r=[]
    dr_rmc=[]
    dr_exp=[]
    error_dr=[]
    Q, sq_rmc, sq_exp, error_sq = read_sq(file2) 
    r, dr_rmc, dr_exp, error_dr = read_dr(file1)
    
    plt.figure()
    plt.subplot(1,2,2)
    plt.title('RMC scattering data refinement',weight = 'bold')
    plt.plot(Q, sq_exp,'k.',label = 'EXP')
    plt.plot(Q, sq_rmc, 'r-', label = 'RMC',linewidth=1.0)
    plt.plot(Q, error_sq, label = 'error',linewidth=0.5)
    plt.xlabel('Q  ( $\AA^{-1}$ )')
    plt.ylabel('S(Q)')  
    plt.legend()

    plt.subplot(1,2,1)
    plt.title('RMC PDF refinement',weight = 'bold')
    plt.plot(r, dr_exp,'k.',label = 'EXP')
    plt.plot(r, dr_rmc, 'r-', label = 'RMC',linewidth=1.0)
    plt.plot(r, error_dr, label = 'error',linewidth=0.5)
    plt.xlabel('r ( $\AA$ )')
    plt.ylabel('D(r)')  
    plt.legend()
    
    plt.tight_layout()
